/*
 * H264RtpPackager.cpp
 *
 *  Created on: 2015年6月2日
 *      Author: chuanjiang.zh@qq.com
 */

#include "H264RtpPackager.h"
#include "H264NaluParser.h"
#include "RtpHeader.h"


namespace av
{

H264RtpPackager::H264RtpPackager():
    m_lastTime()
{

}

H264RtpPackager::~H264RtpPackager()
{
}

void H264RtpPackager::slice(const MediaPacket& pkt, int maxSize, RtpPackagerSink* pSink)
{
    if (!pSink)
    {
        return;
    }

    uint8_t* data = pkt.data;
    int size = pkt.size;

    int count = 0;
    if (H264NaluParser::startWithH264Code(data, size))
    {
        size_t begin = 0;
        bool found = true;

        while (found)
        {
            size_t pos = H264NaluParser::findH264StartCode(data, size, begin + 3);
            if (pos != (size_t)-1)
            {
                sliceNaluPtr(data + begin, data + pos, (uint32_t)pkt.pts, false, maxSize, pSink);
                begin = pos;
                count ++;
            }
            else
            {
                break;
            }
        }

        sliceNaluPtr(data + begin, data + size, (uint32_t)pkt.pts, true, maxSize, pSink);
        count ++;
    }
    else
    {
        sliceNalu(data, size, (uint32_t)pkt.pts, true, maxSize, pSink);
        count ++;
    }
}

void H264RtpPackager::sliceNaluPtr(uint8_t* data, uint8_t* end, uint32_t timestamp, bool mark,
        int maxSize, RtpPackagerSink* pSink)
{
    if (end <= data)
    {
        return;
    }

    size_t size = (end - data);
    NaluPacket nalu;
    if (H264NaluParser::parseNalu(data, size, nalu))
    {
        sliceNalu(data + nalu.prefix, size - nalu.prefix, timestamp, mark, maxSize, pSink);
    }
    else
    {
        sliceNalu(data, size, timestamp, mark, maxSize, pSink);
    }
}

void H264RtpPackager::sliceNalu(uint8_t* data, int size, uint32_t timestamp, bool mark,
               int maxSize, RtpPackagerSink* pSink)
{
    if (size <= maxSize)
    {
        RtpPacket pkt;
        pkt.ts = timestamp;
        pkt.mark = mark;
        pkt.data = data;
        pkt.size = size;
		pkt.pt = m_payload;

        pSink->onSlicePacket(pkt);
        return;
    }

    uint8_t naluType = data[0];
    maxSize -= 2;

    data += 1;
    size -= 1;

    RtpPacket pkt;
    pkt.ts = timestamp;
    pkt.mark = false;
	pkt.pt = m_payload;

    bool start = true;
    bool end = false;
    while (size > maxSize)
    {
        m_sliceBuffer.clear();
        m_sliceBuffer.write(naluType);
        m_sliceBuffer.write((uint8_t)0);
        m_sliceBuffer.write(data, maxSize);

        pkt.data = m_sliceBuffer.data();
        pkt.size = m_sliceBuffer.size();

        pkt.data[0] = (naluType & 0xE0) | 28;

        FU_HEADER* pFuHeader = (FU_HEADER*)(pkt.data + 1);
        pFuHeader->TYPE = naluType & 0x1F;
        pFuHeader->S            = start;
        pFuHeader->E            = end;
        pFuHeader->R            = 0;

        pSink->onSlicePacket(pkt);

        data += maxSize;
        size -= maxSize;

        start = false;
    }

    // last packet
    pkt.mark = mark;

    end = true;

    m_sliceBuffer.clear();
    m_sliceBuffer.write(naluType);
    m_sliceBuffer.write((uint8_t)0);
    m_sliceBuffer.write(data, size);

    pkt.data = m_sliceBuffer.data();
    pkt.size = m_sliceBuffer.size();

    pkt.data[0] = (naluType & 0xE0) | 28;

    FU_HEADER* pFuHeader = (FU_HEADER*)(pkt.data + 1);
    pFuHeader->TYPE = naluType & 0x1F;
    pFuHeader->S            = start;
    pFuHeader->E            = end;
    pFuHeader->R            = 0;

    pSink->onSlicePacket(pkt);
}

bool H264RtpPackager::join(const RtpPacket& pktIn, MediaPacket& pktOut)
{
    static unsigned char s_startCode[] = { 0, 0, 0, 1};

    uint8_t* data = pktIn.data;
    int length = pktIn.size;

    uint8_t naluType = data[0] & 0x1F;
    bool fullPacket = false;
    bool keyFrame = (naluType == 5);
    if (naluType == 5)
    {
        keyFrame = true;
    }
    else if (naluType == 7)
    {
        keyFrame = false;
    }
    else if (naluType == 8)
    {
        keyFrame = false;
    }

    if (naluType == 28 || naluType == 29)
    {
        unsigned char startBit = data[1]&0x80;
        unsigned char endBit = data[1]&0x40;

        if( startBit )
        {
            naluType = ( data[0] & 0xE0 ) | ( data[1] & 0x1F );
            m_joinBuffer.clear();
            m_joinBuffer.write(s_startCode, sizeof(s_startCode));
            m_joinBuffer.write(&naluType, 1);
        }
        else if (m_joinBuffer.empty())
        {
            naluType = ( data[0] & 0xE0 ) | ( data[1] & 0x1F );
            m_joinBuffer.clear();
            m_joinBuffer.write(s_startCode, sizeof(s_startCode));
            m_joinBuffer.write(&naluType, 1);
        }

        data += 2;
        length -= 2;

        m_joinBuffer.write(data, length);

        if (endBit)
        {
            fullPacket = true;
            naluType = (m_joinBuffer[sizeof(s_startCode)] & 0x1F);
            keyFrame = (naluType == 5);
        }
    }
    else
    {
        if (m_joinBuffer.empty())
        {
            m_joinBuffer.write(s_startCode, sizeof(s_startCode));
        }
        else
        {
            //size_t len = m_joinBuffer.length();
        }

        m_joinBuffer.write(data, length);
        fullPacket = true;
    }

    if (fullPacket)
    {
    	pktOut.write(m_joinBuffer.data(), m_joinBuffer.length());
        pktOut.type = MEDIA_TYPE_VIDEO;
        pktOut.pts = pktIn.ts;
        pktOut.duration = 0;

        //pktOut.duration = (int)(pktIn.ts - m_lastTime);
        //if (pktOut.duration < 0)
        //{
        //    pktOut.duration = 0;
        //}

        if (keyFrame)
        {
        	pktOut.flags = (pktOut.flags | MEDIA_FLAG_KEY);
        }

        if (naluType == 7)
        {
        	pktOut.flags = (pktOut.flags | MEDIA_FLAG_KEY);
        }
        else if (naluType == 8)
        {
        	pktOut.flags = (pktOut.flags | MEDIA_FLAG_KEY);
        }

        m_joinBuffer.clear();
        m_lastTime = pktIn.ts;
    }

    return fullPacket;
}


void H264RtpPackager::reset()
{
    m_sliceBuffer.clear();

    m_joinBuffer.clear();
    m_lastTime = 0;
}

bool H264RtpPackager::getFormat(MediaFormat& fmt)
{
	fmt.m_codec = MEDIA_CODEC_H264;
	return true;
}


}
